#!/usr/bin/env python

import os
import sys
import struct

CUBE_FORMAT = '>BBQfffffffB'

class Cube:
    def __init__(self,data):
        lv,id,mx,my,mz,Mx,My,Mz,axis,split,skinMask,child = struct.unpack('>BQffffffifBB',data)
        self.lv = lv
        self.id = id
        self.minPts = ( mx, my, mz )
        self.maxPts = ( Mx, My, Mz )
        self.axis = axis
        self.split = split
        self.skinMask = skinMask
        self.childFlag = child
        self.left = None
        self.right = None
        self.key = ( self.lv, self.id )
        self.parent = None
        self.sample = []
        self.marked = False

    def isLeaf(self):
        return self.childFlag == 0

    def isRoot(self):
        return self.parent is None

    def hasLeft(self):
        return ( self.childFlag & 0x01 ) > 0

    def hasRight(self):
        return ( self.childFlag & 0x02 ) > 0

    def contains(self,x,y,z):
        if self.minPts[0] > x: return False
        if self.minPts[1] > y: return False
        if self.minPts[2] > z: return False
        if self.maxPts[0] < x: return False
        if self.maxPts[1] < y: return False
        if self.maxPts[2] < z: return False
        return True

    def lookup(self,x,y,z):
        if self.isLeaf():
            return self
        elif self.left.contains(x,y,z):
            return self.left.lookup(x,y,z)
        else:
            return self.right.lookup(x,y,z)

    def volume(self):
        return (self.maxPts[0] - self.minPts[0])*(self.maxPts[1] - self.minPts[1])*(self.maxPts[2] - self.minPts[2])

    def mark(self):
        node = self
        while node is not None and not node.marked:
            node.marked = True
            node = node.parent

    def isMarked(self):
        return self.marked

    def delete(self):
        if self.parent is None:
            return

        if self is self.parent.left:
            self.parent.left = None
            self.parent.childFlag -= 1
        else:
            self.parent.right = None
            self.parent.childFlag -= 2
    
    def _collectSample(self,s):
        if len(self.sample):
            s.extend(self.sample)
        if self.left is not None:
            self.left._collectSample(s)
        if self.right is not None:
            self.right._collectSample(s)


    def splitNode(self,points = None):
        import copy
        import pdb

        if points is None:
            if not len(self.sample):
                # collect samples
                self._collectSample(self.sample)
            points = self.sample

        axis = self.lv % 3
        points.sort(key = lambda x: x[axis])
        l = len(points)
        mid = l/2

        #if mid == 0 or axis < 0 or axis > 2:
        #    pdb.set_trace()

        if l % 2 == 0:
            self.split = (points[mid][axis] + points[mid-1][axis]) * 0.5
        else:
            self.split = points[mid][axis]

        self.childFlag = 3

        l = copy.copy(self)

        l.lv = self.lv+1
        l.id = self.id * 2
        l.split = 0.0
        l.childFlag = 0
        l.key = ( l.lv, l.id )
        l.parent = self
        l.minPts = copy.copy(self.minPts)

        r = copy.copy(l)
        r.id = self.id * 2 + 1
        r.key = ( r.lv, r.id )
        r.maxPts = copy.copy(self.maxPts)

        lM = list(self.maxPts)
        lM[axis] = self.split
        l.maxPts = tuple(lM)

        rm = list(self.minPts)
        rm[axis] = self.split
        r.minPts = tuple(rm)

        self.left = l
        self.right = r

        self.left.sample = points[:mid]
        self.right.sample = points[mid:]

        # sanity check
        minSide = min(self.split - l.minPts[axis],r.maxPts[axis] - self.split)
        if minSide < 2*0.00026042:
            print >> sys.stderr, "Too small split at cube %s (%f);[%s;%s]" % (self,minSide,self.left,self.right)

        return ( self.left, self.right )

    def splitUntil(self,numSplit = 1):
        newNodes = []
        leafs = [ self ]
        for i in range(0,numSplit):
            newLeafs = []
            for n in leafs:
                l,r = n.splitNode()
                newNodes.append(l)
                newNodes.append(r)
                newLeafs.append(l)
                newLeafs.append(r)
            leafs = newLeafs
        return newNodes

    def isLoaded(self):
        if ( self.childFlag & 0x01 ) != 0 and self.left is None:
            return False
        if ( self.childFlag & 0x02 ) != 0 and self.right is None:
            return False
        return True

    def addChild(self,child):
        if self.hasLeft() and self.left is None:
            self.left = child
            child.parent = self
        elif self.hasRight() and self.right is None:
            self.right = child
            child.parent = self
        else:
            raise RuntimeError()

    def dropChildren(self):
        self.left = None
        self.right = None
        self.childFlag = 0

    def pack(self,out):
        if self.left is None and (self.childFlag & 0x01) > 0:
            print >> sys.stderr, "Inconsistent node at cube (L-None,1) %s" % self
        if self.left is not None and (self.childFlag & 0x01) == 0:
            print >> sys.stderr, "Inconsistent node at cube (L-CUBE,0) %s" % self
        if self.right is None and (self.childFlag & 0x02) > 0:
            print >> sys.stderr, "Inconsistent node at cube (R-None,2) %s" % self
        if self.right is not None and (self.childFlag & 0x02) == 0:
            print >> sys.stderr, "Inconsistent node at cube (R-CUBE,0) %s" % self
        if self.childFlag == 1 or self.childFlag == 2:
            print >> sys.stderr, "Inconsistent node at cube (single child) %s" % self

        out.write(struct.pack('>BQ',self.lv,self.id))
        out.write(struct.pack('>fff',*self.minPts))
        out.write(struct.pack('>fff',*self.maxPts))
        out.write(struct.pack('>fB',self.split,self.childFlag))
        if self.hasLeft():
            self.left.pack(out)
        if self.hasRight():
            self.right.pack(out)

    def __str__(self):
        return "%d;%d;%s;%s;%d;%f;%f" % ( self.lv, self.id, self.minPts, self.maxPts, self.axis, self.split, self.volume() )

    def dump(self,out = sys.stdout,leafOnly = False):
        if self.hasLeft():
            self.left.dump(out,leafOnly)
        if self.hasRight():
            self.right.dump(out,leafOnly)
        if not leafOnly or (leafOnly and self.isLeaf()):
            print >> out, self

def loadCube(fn):
    recSize = struct.calcsize('>BQffffffifBB')
    input = open(fn,'rb',4*1024*1024)
    root = None
    stack = []
    nodeMap = {}

    while True:
        data = input.read(recSize)
        if not data:
            break

        c = Cube(data)
        nodeMap[c.key] = c

        if len(stack) > 0:
            stack[-1].addChild(c)

        if c.isLeaf():
            # roll up stack until we find non-loaded node
            node = stack[-1]
            while node is not None and node.isLoaded():
                stack.pop()
                if len(stack):
                    node = stack[-1]
                else:
                    node = None
        else:
            if root is None:
                root = c
            stack.append(c)
    input.close()

    return ( root, nodeMap )

def loadSample(fn):
    points = []
    input = open(fn,'rb',4*1024*1024)
    while True:
        data = input.read(20)
        if not data:
            break
        id,x,y,z = struct.unpack('>Qfff',data)
        points.append( (x,y,z) )
    input.close()
    return points

def splitAndMark(nodeMap,node,numSplit):
    for n in node.splitUntil(numSplit):
        nodeMap[n.key] = n
        n.mark()

if __name__ == '__main__':
    from optparse import OptionParser
    
    if len(sys.argv) < 3:
        print >> sys.stderr, "Usage: %s command args..." % sys.argv[0]
        sys.exit(-1)

    cmd = sys.argv[1]

    parser = OptionParser()
    parser.add_option('-o','--out',dest='output',default=None,help='output file')
    parser.add_option('-s','--sample',dest='sample',default=None,help='sample file')
    parser.add_option('-l','--leafOnly',dest='leafOnly',default=False,action='store_true',help='filter non-leaf nodes')
    parser.add_option('-d','--depth',dest='depth',default=0,type='int',help='desired depth of tree')

    ( options, args ) = parser.parse_args(sys.argv[2:])

    if len(args) == 0:
        parser.error("No input partition file")

    inputFiles = []
    for input in args:
        root, nodeMap = loadCube(input)
        inputFiles.append( ( input, root, nodeMap ) )

    sample = []
    if options.sample is not None:
        sample = loadSample(options.sample)

    output = sys.stdout
    if options.output is not None:
        output = open(options.output,'w',4*1024*1024)

    if cmd == 'dump':
        for fn, root, nodeMap in inputFiles:
            print >> output, "File: %s" % fn
            root.dump(output,options.leafOnly)
    elif cmd == 'split':
        fn, root, nodeMap = inputFiles[0]
        # distribute all sample
        for s in sample:
            n = root.lookup(*s)
            n.sample.append(s)

        # now split
        currentLv = max(map(lambda x: x[0],nodeMap.keys()))
        if currentLv < options.depth:
            # do split
            leafs = [ x for x in nodeMap.values() if x.isLeaf() ]
            while currentLv < options.depth:
                newLeafs = []
                for n in leafs:
                    l,r = n.splitNode()
                    nodeMap[l.key] = l
                    nodeMap[r.key] = r
                    newLeafs.append(l)
                    newLeafs.append(r)
                leafs = newLeafs
                currentLv += 1
            root.pack(output)
        elif currentLv > options.depth:
            # do prune
            newLeafs = [ x for x in nodeMap.keys() if x[0] == options.depth ]
            for k in newLeafs:
                l = nodeMap[k]
                l.dropChildren()
            root.pack(output)
    elif cmd == 'edit':
        fn, root, nodeMap = inputFiles[0]
        # distribute all sample
        for s in sample:
            n = root.lookup(*s)
            n.sample.append(s)

        # mark all level-7 nodes
        for n in [ x for x in nodeMap.keys() if x[0] == 7 ]:
            nodeMap[n].mark()

        for n in nodeMap[(7,39)].splitUntil(1):
            nodeMap[n.key] = n
            n.mark()

        splitAndMark(nodeMap,nodeMap[(7,5)],1)
        splitAndMark(nodeMap,nodeMap[(7,116)],1)
        splitAndMark(nodeMap,nodeMap[(7,62)],1)
        splitAndMark(nodeMap,nodeMap[(7,58)],1)

        splitAndMark(nodeMap,nodeMap[(10,954)],4)
        splitAndMark(nodeMap,nodeMap[(10,284)],3)
        splitAndMark(nodeMap,nodeMap[(10,285)],3)
        splitAndMark(nodeMap,nodeMap[(10,3)],3)
        splitAndMark(nodeMap,nodeMap[(10,340)],3)
        splitAndMark(nodeMap,nodeMap[(10,266)],3)
        splitAndMark(nodeMap,nodeMap[(10,947)],2)
        splitAndMark(nodeMap,nodeMap[(10,739)],2)
        splitAndMark(nodeMap,nodeMap[(10,147)],2)
        splitAndMark(nodeMap,nodeMap[(10,615)],1)
        splitAndMark(nodeMap,nodeMap[(10,270)],1)
        splitAndMark(nodeMap,nodeMap[(10,429)],1)
        splitAndMark(nodeMap,nodeMap[(10,476)],1)
        splitAndMark(nodeMap,nodeMap[(10,952)],1)

        toRemove = [ x for x in nodeMap.values() if x.lv > 7 and not x.isMarked() and ( (x.parent.lv == 7 and x.parent.childFlag == 0) or not x.parent.isMarked() ) ]

        for n in toRemove:
            n.delete()

        root.pack(output)

    if options.output is not None:
        output.close()

#    root, nodeMap = loadCube('dense-8.pf')
#   root2, nodeMap2 = loadCube('dense-8-manual.pf')
#   node238 = nodeMap2[(8,238)]
#   node238.childFlag = 3
#   node476 = nodeMap[(9,476)]
#   node477 = nodeMap[(9,477)]
#   node238.addChild(node476)
#   node238.addChild(node477)
#   node476.updateIlv()
#    out = open('dense-8-manual-new.pf','w',65536)
#    root2.pack(out)
#    out.close()
#    root2.dump()
#    root, nodeMap = loadCube('dense-8-manual-new.pf')
#    samples = loadSample('sample.8')
#    node477 = nodeMap[(9,477)]
#    print node477
#    points = [ p for p in samples if node477.contains(*p) ]
#    node477.splitNode(points)
#
#    #root.dump()
#    out = open('dense-8-manual-new2.pf','w',65536)
#    root.pack(out)
#    out.close()
